package com.meilai.project.core.excel;

import org.apache.poi.xssf.streaming.SXSSFWorkbook;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.io.OutputStream;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author xchen
 * @since 2022-02-16
 */
public class ExcelGenerator {

	private SXSSFWorkbook workBook;

	private Map<String, AbstractSheetGenerator> sheetGeneratorMap;

	public static ExcelGenerator createExcel() {
		ExcelGenerator generator = new ExcelGenerator();
		generator.workBook = new SXSSFWorkbook();
		generator.sheetGeneratorMap = new HashMap<>();
		return generator;
	}

	public ExcelGenerator sheet(ReportSheet reportSheet) {
		sheet(reportSheet, new DefaultSheetGenerator());
		return this;
	}

	public ExcelGenerator sheet(ReportSheet reportSheet, AbstractSheetGenerator sheetGenerator) {
		reportSheet.init();
		sheetGenerator.init(this.workBook, reportSheet);
		sheetGeneratorMap.put(reportSheet.getSheetName(), sheetGenerator);
		return this;
	}

	public ExcelGenerator write(List data) {
		if (sheetGeneratorMap.size() != 1) {
			throw new ExportException("sheet超过一个，必须指定sheetName");
		}
		String sheetName = null;
		for (String sheet : sheetGeneratorMap.keySet()) {
			sheetName = sheet;
			break;
		}
		return write(sheetName, data);

	}

	public ExcelGenerator write(String sheetName, List data) {
		AbstractSheetGenerator abstractSheetGenerator = sheetGeneratorMap.get(sheetName);
		if (null == abstractSheetGenerator) {
			throw new ExportException(sheetName + "不存在！");
		}
		abstractSheetGenerator.process(data);
		return this;
	}

	public void save(HttpServletResponse response, String fileName) {
		try {
			for (Map.Entry<String, AbstractSheetGenerator> sheetGeneratorEntry : sheetGeneratorMap.entrySet()) {
				sheetGeneratorEntry.getValue().afterProcess();
			}
			response.setContentType("application/octet-stream");
			response.setHeader("Content-Disposition", "attachment; filename=" + fileName + ".xls");
			ServletOutputStream outputStream = response.getOutputStream();
			this.workBook.write(outputStream);
			outputStream.flush();
		} catch (Exception e) {
			throw new ExportException(e);
		}
	}

	public void save(OutputStream outputStream) {
		try {
			for (Map.Entry<String, AbstractSheetGenerator> sheetGeneratorEntry : sheetGeneratorMap.entrySet()) {
				sheetGeneratorEntry.getValue().afterProcess();
			}
			this.workBook.write(outputStream);
			outputStream.flush();
		} catch (Exception e) {
			throw new ExportException(e);
		}
	}


}
